# Basic Code for embedding.

import jax
import jax.numpy as jnp
import numpy as np
import haiku as hk

from typing import Optional, Union, Callable
from jax import nn as nn
from ..model.interaction.molct_interaction import InteractionUnit
from ..utils.rbf import get_rbf
from ..utils.cutoff import get_cutoff
from ..utils.filter import get_filter
from ..utils.activation import get_activation

class Embedding(hk.Module):
    r"""Embedding Base.

    ### Args:

        fp_type:   Float type. Default: jnp.float32.
        
        int_type:   Int type. Default: jnp.int32.

        dim_node (int):   Node embedding dimension.

        dim_edge (int):   Edge embedding dimension.

        activation (str, Callable):  Activation function. Default: 'silu'.

        name (str):   Module name. Default: 'embedding'.
    """

    def __init__(self,
                 fp_type,
                 int_type,
                 dim_node: int,
                 dim_edge: int,
                 activation: Union[str, Callable] = 'silu',
                 name: str = 'embedding',):
        super().__init__(name=name)

        self.fp_type = fp_type
        self.int_type = int_type

        self.dim_node = int(dim_node)
        self.dim_edge = int(dim_edge)
        self.activation = get_activation(activation)

    def __call__(self, *args, **kwargs):
        raise NotImplementedError